#include <Windows.h>
#include <iostream>
#include <unistd.h>
#include <stdlib.h>

enum class CmdColors
{
    CLEAR = 0,
    RED,
    GREEN,
    BLUE,
    YELLOW,
    PURPLE,
    BLACK
};

std::ostream &operator<<(std::ostream &os, const CmdColors c)
{
    HANDLE hConsole = GetStdHandle(STD_OUTPUT_HANDLE);
    switch (c)
    {
    case CmdColors::CLEAR:
        SetConsoleTextAttribute(hConsole, FOREGROUND_INTENSITY | FOREGROUND_RED | FOREGROUND_GREEN | FOREGROUND_BLUE);
        break;
    case CmdColors::RED:
        SetConsoleTextAttribute(hConsole, FOREGROUND_RED | FOREGROUND_INTENSITY);
        break;
    case CmdColors::GREEN:
        SetConsoleTextAttribute(hConsole, FOREGROUND_GREEN | FOREGROUND_INTENSITY);
        break;
    case CmdColors::BLUE:
        SetConsoleTextAttribute(hConsole, FOREGROUND_BLUE | FOREGROUND_INTENSITY);
        break;
    case CmdColors::YELLOW:
        SetConsoleTextAttribute(hConsole, FOREGROUND_RED | FOREGROUND_GREEN | FOREGROUND_INTENSITY);
        break;
    case CmdColors::PURPLE:
        SetConsoleTextAttribute(hConsole, FOREGROUND_RED | FOREGROUND_BLUE | FOREGROUND_INTENSITY);
        break;
    case CmdColors::BLACK:
        SetConsoleTextAttribute(hConsole, FOREGROUND_RED | FOREGROUND_GREEN | FOREGROUND_BLUE);
        break;
    default:
        break;
    }
    return os;
}

void print_error(int current_time = -1)
{
    using namespace std;
    cout << CmdColors::RED;
    if (current_time >= 0)
    {
        cout << "(" << current_time << ")";
    }
    cout << "Failed, retrying..."
         << CmdColors::CLEAR
         << endl;
}

void run_command(const char *cmd, int retry_time, int exit_code=0)
{
    using namespace std;
    bool success = false;
    int try_time = 0;
    while (!success)
    {
        try_time++;
        success = system(cmd) == exit_code;
        // this time, ok!
        if (success)
        {
            break;
        }
        // not ok!
        else
        {
            // no retry time limit
            if (retry_time < 0)
            {
                print_error(try_time);
                continue;
            }
            // has time left
            if (retry_time > 0)
            {
                print_error(try_time);
                retry_time--;
            }
            // no time left
            else
            {
                cout << CmdColors::RED
                     << "Failed after " << try_time << " times, exiting..."
                     << CmdColors::CLEAR
                     << endl;
                break;
            }
        }
    }
}

bool is_number(const char* str)
{
    for (int i = 0; str[i] != '\0'; i++)
    {
        if (!isdigit(str[i]))
        {
            return false;
        }
    }
    return true;
}

int main(int argc, char *argv[])
{
    using namespace std;

    if (argc != 4 && argc != 6)
    {
        cout << CmdColors::RED
             << "Usage: "
             << CmdColors::CLEAR
             << "until_success <command> "
             << CmdColors::GREEN
             << "-t <times> -c <exit_code>"
             << CmdColors::CLEAR
             << endl;
        cout << "Use \"\" to wrap your command if it contains spaces." << endl;
        return 1;
    }
    // extract args[1] [2] ...
    char *except_first_args[argc - 1];
    for (int i = 1; i < argc; i++)
    {
        except_first_args[i - 1] = argv[i];
    }
    const char *optstring = "t:c:";
    int opt;
    int max_retry_time = 20;
    int exit_code = 0;

    while((opt = getopt(argc, except_first_args, optstring)) != -1)
    {
        switch (opt)
        {
        case 't':
            if (is_number(optarg))
            {
                max_retry_time = atoi(optarg);
            }
            break;
        case 'c':
            if (is_number(optarg))
            {
                exit_code = atoi(optarg);
            }
            break;
        default:
            cout << "unknown opt: " << opt << endl;
            break;
        }
    }
    cout << "times:     " << max_retry_time << endl;
    cout << "exit code: " << exit_code << endl;
    run_command(argv[1], max_retry_time, exit_code);

    return 0;
}